import torch
from torch.nn import Module


class FocalLoss(Module):

    def __init__(self, alpha=0.25, gamma=2, eps=1e-12, reduction='mean', **kwargs):
        super(FocalLoss, self).__init__(**kwargs)
        self.alpha = alpha
        self.gamma = gamma
        self.eps = eps
        self.reduction = reduction

    def forward(self, input, target):
        input = input.sigmoid()
        onehot = input.new_zeros(input.size())
        onehot.scatter_(dim=1, index=target.unsqueeze(1), value=1)

        aux = torch.ones_like(onehot)
        mask = (onehot.long() == 1)

        alpha = torch.where(mask, self.alpha * aux, (1 - self.alpha) * aux)
        pt = torch.where(mask, input, 1 - input)

        loss = - alpha * ((1 - pt) ** self.gamma) * torch.log(torch.min(pt + self.eps, aux))
        loss = torch.sum(loss, dim=-1)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        elif self.reduction == 'none':
            return loss
